//
// Created by ZhaoXiaoFei on 2022/8/2.
//

#ifndef VINS_ESTIMATOR_MARGINALIZATION_FACTOR_HPP
#define VINS_ESTIMATOR_MARGINALIZATION_FACTOR_HPP
#include <ceres/ceres.h>
#include <glog/logging.h>

struct ResidualBlockInfo{
public:
    ResidualBlockInfo(ceres::CostFunction* costFunction_, ceres::LossFunction* lossFunction_,
                      const std::vector<double *> &parameter_blocks_, const std::vector<int> &drop_set_)
                      : costFunction(costFunction_), lossFunction(lossFunction_),
                        parameter_blocks(parameter_blocks_), drop_set(drop_set_){
    }

    void Evaluate(){
//        LOG(INFO) << "1111111111111111" << std::endl;
        residuals.resize( costFunction->num_residuals() );

        std::vector<int> block_sizes = costFunction->parameter_block_sizes();
        raw_jacobians = new double* [block_sizes.size()];
        jacobians.resize(block_sizes.size());

//        LOG(INFO) << "costFunction->num_residuals(): " << costFunction->num_residuals() << std::endl;
        for(int i = 0; i < block_sizes.size(); i++){
//            LOG(INFO) << "block_sizes[i]: " << block_sizes[i] << std::endl;
            jacobians[i].resize(costFunction->num_residuals(), block_sizes[i]);
//            LOG(INFO) << "1111111111111111" << std::endl;
            raw_jacobians[i] = jacobians[i].data();
        }
//        LOG(INFO) << "1111111111111111" << std::endl;
        costFunction->Evaluate(parameter_blocks.data(), residuals.data(), raw_jacobians);
//        LOG(INFO) << "1111111111111111" << std::endl;
        if(lossFunction){
            double residuals_scal_, alpha_sq_norm_;
            double s, rho[3];
            s = residuals.squaredNorm();
            lossFunction->Evaluate(s, rho);
            LOG(INFO) << "sq_norm: " << s << " rho: " << rho[0]  << " "
                                                            << rho[1] << " " << rho[2] << std::endl;

            double sqrt_rho1 = std::sqrt(rho[1]);

            if(s == 0.0 || rho[2] <= 0.0){
                residuals_scal_ = sqrt_rho1;
                alpha_sq_norm_ = 0.0;
            }
            else{
                const double D = 1.0 + 2.0 * s * rho[2] / rho[1];
                const double  alpha = 1.0 - std::sqrt(D);
                residuals_scal_ = sqrt_rho1 / (1 - alpha);
                alpha_sq_norm_ = alpha / s;
            }

            for(int i = 0; i < parameter_blocks.size(); i++){
                jacobians[i] = sqrt_rho1 * (jacobians[i] - alpha_sq_norm_ * residuals * residuals.transpose() * jacobians[i]);
            }

            residuals = residuals_scal_ * residuals;

        }
    }

    ceres::CostFunction* costFunction = nullptr;
    ceres::LossFunction* lossFunction = nullptr;
    std::vector<double*> parameter_blocks;
    std::vector<int> drop_set;

    double** raw_jacobians;
    std::vector<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> jacobians;
    Eigen::VectorXd residuals;
};

class MarginalizationInfo{
public:
    MarginalizationInfo(){}
    inline int localSize(int globalSize){
        return globalSize == 7 ? 6 : globalSize;
    }
    inline int globalSize(int localSize){
        return localSize == 6 ? 7 : localSize;
    }
    void addResidualBlockInfo(ResidualBlockInfo* residualBlockInfo){
        factors.emplace_back(residualBlockInfo);

        std::vector<double*>& parameter_blocks = residualBlockInfo->parameter_blocks;
        std::vector<int32_t> parameter_block_sizes = residualBlockInfo->costFunction->parameter_block_sizes();

//        LOG(INFO) << "parameter_block_sizes.size: " << parameter_block_sizes.size() << std::endl;
//        LOG(INFO) << "residualBlockInfo->parameter_blocks.size(): " << residualBlockInfo->parameter_blocks.size() << std::endl;

        //所有的参数size。
        for(int i = 0; i < residualBlockInfo->parameter_blocks.size(); i++){
            double* addr = parameter_blocks[i];
            int size = parameter_block_sizes[i];
            parameter_block_size[reinterpret_cast<long>(addr)] = size;
        }
        //首先对于需要抛弃的参数。
        for(int i = 0; i < residualBlockInfo->drop_set.size(); i++){
            double* addr = parameter_blocks[residualBlockInfo->drop_set[i]];
            parameter_block_idx[reinterpret_cast<long>(addr)] = 0;
        }
    }

    void preMarginaliza(){
        for(auto it : factors){
//            LOG(INFO) << "1111111111111111" << std::endl;
            it->Evaluate();
//            LOG(INFO) << "1111111111111111" << std::endl;
            std::vector<int> parameter_block_sizes = it->costFunction->parameter_block_sizes();
            for(int i = 0; i < parameter_block_sizes.size(); i++){
//                LOG(INFO) << "1111111111111111" << std::endl;
                long addr = reinterpret_cast<long>(it->parameter_blocks[i]);//地址
                int size = parameter_block_sizes[i];
                if(parameter_block_data.find(addr) == parameter_block_data.end()){
                    double* data = new double[size];
                    memcpy(data, it->parameter_blocks[i], sizeof(double ) * size);
                    parameter_block_data[addr] = data;
                }
            }
        }
    }

    void marginalize(){
//        LOG(INFO) << "1111111111111111" << std::endl;
        int pos = 0;
        for(auto& it : parameter_block_idx){
            it.second = pos;
            pos += localSize(parameter_block_size[it.first]);
        }
        m = pos;
        for(auto& it : parameter_block_size){
            if(parameter_block_idx.find(it.first) == parameter_block_idx.end()){
                parameter_block_idx[it.first] = pos;
                pos += localSize(it.second);
            }
        }
        r = pos - m;

        Eigen::MatrixXd H(pos, pos);
        Eigen::VectorXd b(pos);
        H.setZero();
        b.setZero();
//        LOG(INFO) << "1111111111111111" << std::endl;
        //对于每一个因子
        for(auto& it : factors){
            //每一个因子所优化的每一个残差块
            for(int i = 0; i < it->parameter_blocks.size(); i++){
                int idx_i = parameter_block_idx[reinterpret_cast<long>(it->parameter_blocks[i])];
                int size_i = localSize(parameter_block_size[reinterpret_cast<long>(it->parameter_blocks[i])]);
                Eigen::MatrixXd jacobian_i = it->jacobians[i].leftCols(size_i);
                for(int j = i; j < it->parameter_blocks.size(); j++){
                    int idx_j = parameter_block_idx[reinterpret_cast<long>(it->parameter_blocks[j])];
                    int size_j = localSize(parameter_block_size[reinterpret_cast<long>(it->parameter_blocks[j])]);
                    Eigen::MatrixXd jacobian_j = it->jacobians[j].leftCols(size_j);
                    if(i == j){
                        H.block(idx_i, idx_j, size_i, size_j) += jacobian_i.transpose() * jacobian_j;
                    }
                    else{
                        H.block(idx_i, idx_j, size_i, size_j) += jacobian_i.transpose() * jacobian_j;
                        H.block(idx_j, idx_i, size_j, size_i) = H.block(idx_i, idx_j, size_i, size_j).transpose();
                    }
                }
                b.segment(idx_i, size_i) += jacobian_i.transpose() * it->residuals;
            }
        }

        Eigen::MatrixXd Hmm = 0.5 * (H.block(0, 0, m, m) + H.block(0, 0, m, m).transpose());

        Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> saes(Hmm);

        Eigen::MatrixXd Hmm_inv = saes.eigenvectors()
                * Eigen::VectorXd((saes.eigenvalues().array() > eps).select(saes.eigenvalues().array().inverse(), 0)).asDiagonal()
                * saes.eigenvectors().transpose();

        Eigen::MatrixXd Hmr = H.block(0, m, m, r);
        Eigen::MatrixXd Hrm = H.block(m, 0, r, m);
        Eigen::MatrixXd Hrr = H.block(m, m, r, r);

        Eigen::VectorXd bmm = b.segment(0, m);
        Eigen::VectorXd brr = b.segment(m, r);

        H = Hrr - Hrm * Hmm_inv * Hmr;
        b = brr - Hrm * Hmm_inv * bmm;

        Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> saes2(H);
        Eigen::VectorXd S = Eigen::VectorXd((saes2.eigenvalues().array() > eps).select(saes2.eigenvalues().array(), 0));
        Eigen::VectorXd S_inv = Eigen::VectorXd((saes2.eigenvalues().array() > eps).select(saes2.eigenvalues().array().inverse(), 0));
        Eigen::VectorXd S_sqrt = S.cwiseSqrt();
        Eigen::VectorXd S_inv_sqrt = S_inv.cwiseSqrt();

        linearized_jacobians = S.cwiseSqrt().asDiagonal() * saes2.eigenvectors().transpose();
        linearized_residuals = S_inv_sqrt.asDiagonal() * saes2.eigenvectors().transpose() * b;

    }

    std::vector<double*> getParameterBlocks(std::unordered_map<long, double*>& addr_shift){

        std::vector<double*> keep_block_addr;
        keep_block_size.clear();
        keep_block_idx.clear();
        keep_block_data.clear();

        for(const auto& it : parameter_block_idx){
            if(it.second >= m){
                keep_block_size.push_back(parameter_block_size[it.first]);
                keep_block_idx.push_back(parameter_block_idx[it.first]);
                keep_block_data.push_back(parameter_block_data[it.first]);
                keep_block_addr.push_back(addr_shift[it.first]);
            }
        }
        sum_block_size = std::accumulate(std::begin(keep_block_size), std::end(keep_block_size), 0);
        return keep_block_addr;
    }

    ~MarginalizationInfo(){

        for(auto it = parameter_block_data.begin(); it != parameter_block_data.end(); ++it)
            delete[] it->second;

        for(int i = 0; i < (int)factors.size(); i++){
            delete[] factors[i]->raw_jacobians;
            delete factors[i]->costFunction;
            delete factors[i];
        }
    }

    std::vector<ResidualBlockInfo* > factors;
    int m, r;
    int sum_block_size;
    std::unordered_map<long, int> parameter_block_size;
    std::unordered_map<long, int> parameter_block_idx;
    std::unordered_map<long, double*> parameter_block_data;

    //边缘化后剩余的参数;
    std::vector<int> keep_block_size;
    std::vector<int> keep_block_idx;
    std::vector<double*> keep_block_data;

    Eigen::MatrixXd linearized_jacobians;
    Eigen::VectorXd linearized_residuals;
    const double eps = 1e-8;
};

class MarginalizationFactor : public ceres::CostFunction{
public:
    MarginalizationFactor(MarginalizationInfo* marginalizationInfo_) : marginalizationInfo(marginalizationInfo_){
        int cnt = 0;
        for(auto& it : marginalizationInfo->keep_block_size){
            mutable_parameter_block_sizes()->push_back(it);
            cnt += it;
        }
        set_num_residuals(marginalizationInfo->r);
    }
    virtual bool Evaluate(double const* const* parameters,
                          double* residuals,
                          double** jacobians) const{
        int m = marginalizationInfo->m;
        int r = marginalizationInfo->r;
        Eigen::VectorXd dx(r);
        for(int i = 0; i < marginalizationInfo->keep_block_size.size(); i++){
            int idx = marginalizationInfo->keep_block_idx[i] - m;
            int size = marginalizationInfo->keep_block_size[i];

            Eigen::VectorXd x = Eigen::Map<const Eigen::VectorXd>(parameters[i], size);
            Eigen::VectorXd x0 = Eigen::Map<const Eigen::VectorXd>(marginalizationInfo->keep_block_data[i], size);
            if(size != 7){
                dx.segment(idx, size) = x - x0;
            }else{
                dx.segment<3>(idx + 0) = x.head<3>() - x0.head<3>();
                dx.segment<3>(idx + 3) = 2.0 * (Eigen::Quaterniond(x0(6), x0(3), x0(4), x0(5)).inverse()
                        * Eigen::Quaterniond(x(6), x(3), x(4), x(5))).vec();
                if ((Eigen::Quaterniond(x0(6), x0(3), x0(4), x0(5)).inverse()
                     * Eigen::Quaterniond(x(6), x(3), x(4), x(5))).w() < 0)
                {
                    dx.segment<3>(idx + 3) = 2.0 * -(Eigen::Quaterniond(x0(6), x0(3), x0(4), x0(5)).inverse() * Eigen::Quaterniond(x(6), x(3), x(4), x(5))).vec();
                }
            }
        }
        Eigen::Map<Eigen::VectorXd>(residuals, r) = marginalizationInfo->linearized_residuals +  marginalizationInfo->linearized_jacobians * dx;
        if(jacobians){

            for(int i = 0; i < marginalizationInfo->keep_block_size.size(); i++){
                if(jacobians[i]){
                    int size = marginalizationInfo->keep_block_size[i];
                    int localSize = marginalizationInfo->localSize(size);
                    int idx = marginalizationInfo->keep_block_idx[i] - m;
                    Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> jacobian(jacobians[i], r, size);
                    jacobian.setZero();
                    jacobian.leftCols(localSize) = marginalizationInfo->linearized_jacobians.middleCols(idx, localSize);
                }
            }
        }
        return true;
    }

    MarginalizationInfo* marginalizationInfo;
};

#endif //VINS_ESTIMATOR_MARGINALIZATION_FACTOR_HPP
